# !/usr/bin/env python
# -*- coding: utf-8 -*-
from pyspark.ml import Pipeline
from pyspark.sql import SparkSession
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.regression import LinearRegression
from pyspark2pmml import PMMLBuilder
from jpmml_evaluator import make_evaluator
from pyspark.sql.functions import col, log, sqrt, abs, exp
import math

model_path = "../out/house_price.pmml"


def save_model():
    # creating a spark session
    spark = SparkSession.builder.appName("house_price2").master("local[*]") \
        .config("spark.jars.packages", "org.jpmml:pmml-sparkml:2.4.1").getOrCreate()
    # reading csv files
    train = spark.read.csv("../data/house-prices/train.csv", inferSchema=True, header=True)
    train.printSchema()

    first = train.head(1)
    print("first", first[0].asDict())
    # train = train.fillna(0.1)
    # 使用na.drop方法删除含有null值的行
    # train = train.na.drop()

    input_cols = ["MSSubClass", "LotArea", "OverallQual", "OverallCond", "BsmtFinSF1",
                   "BsmtFinSF2", "BsmtUnfSF", "TotalBsmtSF", "1stFlrSF", "2ndFlrSF",
                   "LowQualFinSF", "GrLivArea", "BsmtFullBath", "BsmtHalfBath",
                   "FullBath", "HalfBath", "BedroomAbvGr", "KitchenAbvGr", "TotRmsAbvGrd",
                   "Fireplaces", "YearBuilt",
                   "YearRemodAdd", "GarageCars", "GarageArea", "WoodDeckSF", "OpenPorchSF",
                   "EnclosedPorch", "3SsnPorch", "ScreenPorch", "PoolArea",
                   "MiscVal", "MoSold", "YrSold", "LotShape2"]
    indexer = StringIndexer(inputCol="LotShape", outputCol="LotShape2")
    train = indexer.fit(train).transform(train)

    # Assembler combines all integer and create a vector which is used as input to predict. Here we have only
    # selected columns with data type as integer
    assembler = VectorAssembler(inputCols=input_cols, outputCol="features")

    (train_df, valid_df) = train.randomSplit([0.7, 0.3], seed=42)
    train_df.show(1)

    lr = LinearRegression(featuresCol='features', labelCol='SalePrice', maxIter=10,
                          regParam=0.8, elasticNetParam=0.1)
    pipeline = Pipeline(stages=[assembler, lr])
    model = pipeline.fit(train_df)

    lr_predictions = model.transform(valid_df)
    lr_predictions.select("prediction", "SalePrice", "features").show(5)

    lr_evaluator = RegressionEvaluator(predictionCol="prediction",
                                       labelCol="SalePrice", metricName="r2")
    print("R Squared (R2) on val data = %g" % lr_evaluator.evaluate(lr_predictions))
    pmmlBuilder = PMMLBuilder(spark.sparkContext, train_df, model).putOption(lr, "compact", True)
    pmmlBuilder.buildByteArray()
    pmmlBuilder.buildFile(model_path)
    spark.stop()


def load_model():
    # model_path = "G:/doc/aimodeler/dev/resource/infinity/model.pmml"
    evaluator = make_evaluator(model_path).verify()
    input_fields = evaluator.getInputFields()
    print("\r\nInput fields: " + str([input_field.getName() for input_field in input_fields]))
    target_fields = evaluator.getTargetFields()
    print("\r\nTarget field(s): " + str([target_field.getName() for target_field in target_fields]))
    output_fields = evaluator.getOutputFields()
    print("\r\nOutput fields: " + str([output_field.getName() for output_field in output_fields]))

    args = {'Id': 1, 'MSSubClass': 60, 'MSZoning': 'RL', 'LotFrontage': '65', 'LotArea': 8450, 'Street': 'Pave',
            'Alley': 'NA', 'LotShape': 'Reg', 'LandContour': 'Lvl', 'Utilities': 'AllPub', 'LotConfig': 'Inside',
            'LandSlope': 'Gtl', 'Neighborhood': 'CollgCr', 'Condition1': 'Norm', 'Condition2': 'Norm',
            'BldgType': '1Fam', 'HouseStyle': '2Story', 'OverallQual': 7, 'OverallCond': 5, 'YearBuilt': 2003,
            'YearRemodAdd': 2003, 'RoofStyle': 'Gable', 'RoofMatl': 'CompShg', 'Exterior1st': 'VinylSd',
            'Exterior2nd': 'VinylSd', 'MasVnrType': 'BrkFace', 'MasVnrArea': '196', 'ExterQual': 'Gd',
            'ExterCond': 'TA', 'Foundation': 'PConc', 'BsmtQual': 'Gd', 'BsmtCond': 'TA', 'BsmtExposure': 'No',
            'BsmtFinType1': 'GLQ', 'BsmtFinSF1': 706, 'BsmtFinType2': 'Unf', 'BsmtFinSF2': 0, 'BsmtUnfSF': 150,
            'TotalBsmtSF': 856, 'Heating': 'GasA', 'HeatingQC': 'Ex', 'CentralAir': 'Y', 'Electrical': 'SBrkr',
            '1stFlrSF': 856, '2ndFlrSF': 854, 'LowQualFinSF': 0, 'GrLivArea': 1710, 'BsmtFullBath': 1,
            'BsmtHalfBath': 0, 'FullBath': 2, 'HalfBath': 1, 'BedroomAbvGr': 3, 'KitchenAbvGr': 1, 'KitchenQual': 'Gd',
            'TotRmsAbvGrd': 8, 'Functional': 'Typ', 'Fireplaces': 0, 'FireplaceQu': 'NA', 'GarageType': 'Attchd',
            'GarageYrBlt': '2003', 'GarageFinish': 'RFn', 'GarageCars': 2, 'GarageArea': 548, 'GarageQual': 'TA',
            'GarageCond': 'TA', 'PavedDrive': 'Y', 'WoodDeckSF': 0, 'OpenPorchSF': 61, 'EnclosedPorch': 0,
            '3SsnPorch': 0, 'ScreenPorch': 0, 'PoolArea': 0, 'PoolQC': 'NA', 'Fence': 'NA', 'MiscFeature': 'NA',
            'MiscVal': 0, 'MoSold': 2, 'YrSold': 2008, 'SaleType': 'WD', 'SaleCondition': 'Normal', 'SalePrice': 208500,
            'LotShape2': 0.0}
    result = evaluator.evaluate(args)

    print('\r\n', result)


if __name__ == '__main__':
    # save_model()
    load_model()
